# LaTeX parser for oxyxy
#
# based on the parser.py in mathtex (http://code.google.com/p/mathtex/) which is itself
# based on mathtext.py in matplotlib (http://matplotlib.sourceforge.net/index.html)

from sympy import *
from sympy.printing import latex

sys.path.append('./pyparsing') # a hack because pyparsing is an svn:external checkout
from pyparsing import *

# TeX to Unicode conversion data
from data import tex2uni

import unittest

# Enable packrat parsing, this gives a ~2x speed-up
ParserElement.enablePackrat()

##############################################################################
# PARSER

def Error(msg):
    """
    Helper class to raise parser errors.
    """
    def raise_error(s, loc, toks):
        raise ParseFatalException(msg + "\n" + s)
    
    empty = Empty()
    empty.setParseAction(raise_error)
    return empty

class MathtexParser(object):
    """
    This is the pyparsing-based parser for math expressions.  It
    actually parses full strings *containing* math expressions, in
    that raw text may also appear outside of pairs of ``$``.
    
    The grammar is based directly on that in TeX, though it cuts a few
    corners.
    """
    _binary_operators = set(r'''
      + *
      \pm             \sqcap                   \rhd
      \mp             \sqcup                   \unlhd
      \times          \vee                     \unrhd
      \div            \wedge                   \oplus
      \ast            \setminus                \ominus
      \star           \wr                      \otimes
      \circ           \diamond                 \oslash
      \bullet         \bigtriangleup           \odot
      \cdot           \bigtriangledown         \bigcirc
      \cap            \triangleleft            \dagger
      \cup            \triangleright           \ddagger
      \uplus          \lhd                     \amalg'''.split())
    
    _relation_symbols = set(r'''
      = < > :
      \leq            \geq             \equiv           \models
      \prec           \succ            \sim             \perp
      \preceq         \succeq          \simeq           \mid
      \ll             \gg              \asymp           \parallel
      \subset         \supset          \approx          \bowtie
      \subseteq       \supseteq        \cong            \Join
      \sqsubset       \sqsupset        \neq             \smile
      \sqsubseteq     \sqsupseteq      \doteq           \frown
      \in             \ni              \propto
      \vdash          \dashv           \dots'''.split())
    
    _arrow_symbols = set(r'''
      \leftarrow              \longleftarrow           \uparrow
      \Leftarrow              \Longleftarrow           \Uparrow
      \rightarrow             \longrightarrow          \downarrow
      \Rightarrow             \Longrightarrow          \Downarrow
      \leftrightarrow         \longleftrightarrow      \updownarrow
      \Leftrightarrow         \Longleftrightarrow      \Updownarrow
      \mapsto                 \longmapsto              \nearrow
      \hookleftarrow          \hookrightarrow          \searrow
      \leftharpoonup          \rightharpoonup          \swarrow
      \leftharpoondown        \rightharpoondown        \nwarrow
      \rightleftharpoons      \leadsto'''.split())
    
    #_spaced_symbols = _binary_operators | _relation_symbols | _arrow_symbols
    
    _punctuation_symbols = set(r', ; . ! \ldotp \cdotp'.split())
    
    _overunder_symbols = set(r'''
       \sum \prod \coprod \bigcap \bigcup \bigsqcup \bigvee
       \bigwedge \bigodot \bigotimes \bigoplus \biguplus
       '''.split())
    
    _overunder_functions = set(
        r"lim liminf limsup sup max min".split())
    
    _dropsub_symbols = set(r'''\int \oint'''.split())
    
    _fontnames = set("rm cal it tt sf bf default bb frak circled scr regular".split())
    
    # map function symbols to corresponding SymPy functions
    # all these functions must take only one argument
    fn_map = { 
#        "inf" : None,
#        "liminf" : None,
#        "limsup" : None,
#        "lim" : None,
#        "min" : None,
#        "max" : None,
        
#        "ker" : None,
#        "deg" : None,
#        "lg" : None,
#        "Pr" : None,
#        "det" : None,
#        "arg" : None,
#        "dim" : None,
#        "gcd" : None,
#        "sup" : None,
#        "hom" : None,
        
        "exp" : exp,
        "ln" : ( lambda x : log(x, e) ),
        #"log" : log,  # log takes two arguments
        
        "sin" : sin,
        "cos" : cos,
        #"sec" : sec,  # secant function doesn't seem to be properly implemented in SymPy at the moment
        #"csc" : csc,  # secant function doesn't seem to be properly implemented in SymPy at the moment
        "tan" : tan,
        "cot" : cot,
        
        "arcsin" : asin,
        "arccos" : acos,
        "arctan" : atan,
        
        "sinh" : sinh,
        "cosh" : cosh,
        "tanh" : tanh,
        "coth" : coth }
    
    _function_names = set(fn_map.keys())
    
    _ambiDelim = set(r"""
      | \| / \backslash \uparrow \downarrow \updownarrow \Uparrow
      \Downarrow \Updownarrow .""".split())
    
    _leftDelim = set(r"( [ { < \lfloor \langle \lceil".split())
    
    _rightDelim = set(r") ] } > \rfloor \rangle \rceil".split())
    
    def __init__(self):
        # All forward declarations are here
        font = Forward().setParseAction(self.font).setName("font")
        latexfont = Forward()
        subsuper = Forward().setParseAction(self.subsuperscript).setName("subsuper")
        placeable = Forward().setName("placeable")
        simple = Forward().setName("simple")
        autoDelim = Forward().setParseAction(self.auto_sized_delimiter)
        self._expression = Forward().setParseAction(self.finish).setName("finish")
        
        float        = Regex(r"[-+]?([0-9]+\.?[0-9]*|\.[0-9]+)")
        
        lbrace       = Literal('{').suppress()
        rbrace       = Literal('}').suppress()
        
        lbracket       = Literal('(').suppress()
        rbracket       = Literal(')').suppress()
        
        start_group  = (Optional(latexfont) - lbrace)
        start_group.setParseAction(self.start_group)
        end_group    = rbrace.copy()
        end_group.setParseAction(self.end_group)
        
        bslash       = Literal('\\')
        
        accent       = oneOf(self._accent_map.keys() +
                             list(self._wide_accents))
        
        function     = oneOf(list(self._function_names))
        
        fontname     = oneOf(list(self._fontnames))
        latex2efont  = oneOf(['math' + x for x in self._fontnames])
        
        space        =(FollowedBy(bslash)
                     + oneOf([r'\ ',
                              r'\/',
                              r'\,',
                              r'\;',
                              r'\quad',
                              r'\qquad',
                              r'\!'])
                      ).setParseAction(self.space).setName('space')
        
        customspace  =(Literal(r'\hspace')
                     - (( lbrace
                        - float
                        - rbrace
                       ) | Error(r"Expected \hspace{n}"))
                     ).setParseAction(self.customspace).setName('customspace')
        
        unicode_range = u"\U00000080-\U0001ffff"
        symbol       =(Regex(UR"([a-zA-Z0-9 +\-*/<>=:,.;!'@()\[\]|%s])|(\\[%%${}\[\]_|])" % unicode_range)
                     | (Combine(
                         bslash
                       + oneOf(tex2uni.keys())
                       ) + FollowedBy(Regex("[^a-zA-Z]")))
                     ).setParseAction(self.symbol).leaveWhitespace()
        
        c_over_c     =(Suppress(bslash)
                     + oneOf(self._char_over_chars.keys())
                     ).setParseAction(self.char_over_chars)
        
        accent       = Group(
                         Suppress(bslash)
                       + accent
                       - placeable
                     ).setParseAction(self.accent).setName("accent")
        
        function     =(Suppress(bslash)
                     + function
                     ).setParseAction(self.function).setName("function")
        
        function_app =( function + lbracket + symbol + rbracket
                     ).setParseAction(self.function_app).setName("function_app")
        
        group        = Group(
                         start_group
                       + ZeroOrMore(
                           autoDelim
                         ^ simple)
                       - end_group
                     ).setParseAction(self.group).setName("group")
        
        font        <<(Suppress(bslash)
                     + fontname)
        
        latexfont   <<(Suppress(bslash)
                     + latex2efont)
        
        frac         = Group(
                       Suppress(Literal(r"\frac"))
                     + ((group + group)
                        | Error(r"Expected \frac{num}{den}"))
                     ).setParseAction(self.frac).setName("frac")
        
        stackrel     = Group(
                       Suppress(Literal(r"\stackrel"))
                     + ((group + group)
                        | Error(r"Expected \stackrel{num}{den}"))
                     ).setParseAction(self.stackrel).setName("stackrel")
        
        binom        = Group(
                       Suppress(Literal(r"\binom"))
                     + ((group + group)
                        | Error(r"Expected \binom{num}{den}"))
                     ).setParseAction(self.binom).setName("binom")
        
        ambiDelim    = oneOf(list(self._ambiDelim))
        leftDelim    = oneOf(list(self._leftDelim))
        rightDelim   = oneOf(list(self._rightDelim))
        rightDelimSafe = oneOf(list(self._rightDelim - set(['}'])))
        genfrac      = Group(
                       Suppress(Literal(r"\genfrac"))
                     + ((Suppress(Literal('{')) +
                         oneOf(list(self._ambiDelim | self._leftDelim | set(['']))) +
                         Suppress(Literal('}')) +
                         Suppress(Literal('{')) +
                         oneOf(list(self._ambiDelim |
                                    (self._rightDelim - set(['}'])) |
                                    set(['', r'\}']))) +
                         Suppress(Literal('}')) +
                         Suppress(Literal('{')) +
                         Regex("[0-9]*(\.?[0-9]*)?") +
                         Suppress(Literal('}')) +
                         group + group + group)
                        | Error(r"Expected \genfrac{ldelim}{rdelim}{rulesize}{style}{num}{den}"))
                     ).setParseAction(self.genfrac).setName("genfrac")
        
        sqrt         = Group(
                       Suppress(Literal(r"\sqrt"))
                     + Optional(
                         Suppress(Literal("["))
                       - Regex("[0-9]+")
                       - Suppress(Literal("]")),
                         default = None
                       )
                     + (group | Error("Expected \sqrt{value}"))
                     ).setParseAction(self.sqrt).setName("sqrt")
        
        operatorname = Group(
                       Suppress(Literal(r"\operatorname"))
                     + ((start_group + Regex("[A-Za-z]+") + end_group)
                        | Error("Expected \operatorname{value}"))
                     ).setParseAction(self.operatorname).setName("operatorname")
        
        placeable   <<(function
                     ^ (c_over_c | symbol)
                     ^ accent
                     ^ group
                     ^ frac
                     ^ stackrel
                     ^ binom
                     ^ genfrac
                     ^ sqrt
                     ^ operatorname
                     )
        
        simple      <<(space
                     | customspace
                     | font
                     | subsuper
                     )
        
        subsuperop   = oneOf(["_", "^"])
        
        subsuper    << Group(
                         ( Optional(placeable)
                         + OneOrMore(
                             subsuperop
                           - placeable
                           )
                         )
                       | placeable
                     )
        
        autoDelim   <<(Suppress(Literal(r"\left"))
                     + ((leftDelim | ambiDelim) | Error("Expected a delimiter"))
                     + Group(
                         autoDelim
                         ^ OneOrMore(simple))
                     + Suppress(Literal(r"\right"))
                     + ((rightDelim | ambiDelim) | Error("Expected a delimiter"))
                     )
        
        #math         = OneOrMore(
        #               autoDelim
        #             ^ simple
        #             ).setParseAction(self.math).setName("math")
        
        math         = OneOrMore(
                       function_app
                     ).setParseAction(self.math).setName("math")
        
        math_delim   = ~bslash + Literal('$')
        
        non_math     = Regex(r"(?:(?:\\[$])|[^$])*"
                     ).setParseAction(self.non_math).setName("non_math").leaveWhitespace()
        
        #self._expression << (
        #    non_math
        #  + ZeroOrMore(
        #        Suppress(math_delim)
        #      + Optional(math)
        #      + (Suppress(math_delim)
        #         | Error("Expected end of math '$'"))
        #      + non_math
        #    )
        #  ) + StringEnd()
        
        self._expression << ( ( Literal("Differentiate") ^ Literal("Integrate") )
            + ( Suppress(math_delim)
            + (math ^ frac ^ sqrt)
            + (Suppress(math_delim)
                 | Error("Expected end of math '$'"))
            ) + Literal("with respect to").suppress()
            + ( Suppress(math_delim)
            + simple
            + (Suppress(math_delim)
                 | Error("Expected end of math '$'"))
            ) ).setParseAction(self.operate).setName("operate")
        
        self.clear()
    
    def clear(self):
        """
        Clear any state before parsing.
        """
        self._expr = None
        self._state_stack = None
        
    
    def parse(self, s):
        """
        Parse expression *s* using the given *fonts_object* for
        output, at the given *fontsize* and *dpi*.
        
        Returns the parse tree of :class:`Node` instances.
        """
        retValue = ""
        self._state_stack = [self.State("foo")]
        try:
            retValue = self._expression.parseString(s)
        except (ParseException, ParseFatalException), err:
            raise ValueError(str(err))
        
        #return self._expr
        return retValue
        
    # The state of the parser is maintained in a stack.  Upon
    # entering and leaving a group { } or math/non-math, the stack
    # is pushed and popped accordingly.  The current state always
    # exists in the top element of the stack.
    class State(object):
        """
        Stores the state of the parser.
        
        States are pushed and popped from a stack as necessary, and
        the "current" state is always at the top of the stack.
        """
        def __init__(self,foo):
            self.foo = foo
            pass
        
        def copy(self):
            return MathtexParser.State(self.foo)
        
    
    def get_state(self):
        """
        Get the current :class:`State` of the parser.
        """
        return self._state_stack[-1]
    
    def pop_state(self):
        """
        Pop a :class:`State` off of the stack.
        """
        self._state_stack.pop()
    
    def push_state(self):
        """
        Push a new :class:`State` onto the stack which is just a copy
        of the current state.
        """
        self._state_stack.append(self.get_state().copy())
    
    # Action handlers
    
    def operate(self, s, loc, toks):
        #print "operate %s,%s" % (loc, toks)
        
        op, exp, wrt = toks[0], toks[1], toks[2]
        
        #print op, exp, wrt
        if op == "Differentiate":
            result = diff(exp, wrt)
        elif op == "Integrate":
            result = integrate(exp, wrt)
        
        #print latex(result)
        return [latex(result)]
    
    def function_app(self, s, loc, toks):
        #print "function_app %s,%s" % (loc, toks)
        
        fn, exp = toks[0],toks[1]
        return fn(exp)
    
    def finish(self, s, loc, toks):
        #print "finish %s,%s" % (loc, toks)
        pass
    
    def math(self, s, loc, toks):
        #print "math %s,%s" % (loc, toks)
        self.pop_state()
    
    def non_math(self, s, loc, toks):
        #print "non_math %s,%s" % (loc, toks)
        s = toks[0].replace(r'\$', '$')
        symbols = [c for c in s]
        # We're going into math now
        self.push_state()
        return [symbols]
    
    def space(self, s, loc, toks):
        print "space %s,%s" % (loc, toks)
    
    def customspace(self, s, loc, toks):
        print "customspace %s,%s" % (loc, toks)
    
    def symbol(self, s, loc, toks):
        #print "symbol %s,%s" % (loc, toks)
        
        symb = toks[0]
        
        if symb == r'(':
            return [symb]
        elif symb == r'(':
            return [symb]
        elif symb == r'+':
            return [symb]
        else:
            try:
                symbol = Symbol(symb.encode('utf-8')) # convert the string from unicode to utf-8
            except ValueError:
                raise ParseFatalException("Unknown symbol: %s" % symb)
            
            return [symbol]
        
    
    _char_over_chars = {
        # The first 2 entires in the tuple are (font, char, sizescale) for
        # the two symbols under and over.  The third element is the space
        # (in multiples of underline height)
        r'AA' : (  ('rm', 'A', 1.0), (None, '\circ', 0.5), 0.0),
    }
    
    def char_over_chars(self, s, loc, toks):
        print "char_over_chars %s,%s" % (loc, toks)
        sym = toks[0]
        state = self.get_state()
        under_desc, over_desc, space = \
            self._char_over_chars.get(sym, (None, None, 0.0))
        if under_desc is None:
            raise ParseFatalException("Error parsing symbol")
    
    _accent_map = {
        r'hat'   : r'\circumflexaccent',
        r'breve' : r'\combiningbreve',
        r'bar'   : r'\combiningoverline',
        r'grave' : r'\combininggraveaccent',
        r'acute' : r'\combiningacuteaccent',
        r'ddot'  : r'\combiningdiaeresis',
        r'tilde' : r'\combiningtilde',
        r'dot'   : r'\combiningdotabove',
        r'vec'   : r'\combiningrightarrowabove',
        r'"'     : r'\combiningdiaeresis',
        r"`"     : r'\combininggraveaccent',
        r"'"     : r'\combiningacuteaccent',
        r'~'     : r'\combiningtilde',
        r'.'     : r'\combiningdotabove',
        r'^'     : r'\circumflexaccent',
        r'overrightarrow' : r'\rightarrow',
        r'overleftarrow'  : r'\leftarrow'
        }
    
    _wide_accents = set(r"widehat widetilde".split())
    
    def accent(self, s, loc, toks):
        print "accent %s,%s" % (loc, toks)
        assert(len(toks)==1)
        state = self.get_state()
        if len(toks[0]) != 2:
            raise ParseFatalException("Error parsing accent")
        accent, sym = toks[0]
    
    def function(self, s, loc, toks):
        #print "function %s,%s" % (loc, toks)
        function_name = toks[0]
        return [self.fn_map[function_name]]
    
    def operatorname(self, s, loc, toks):
        print "operatorname %s,%s" % (loc, toks)
        return self.function(s, loc, toks[0])
    
    def start_group(self, s, loc, toks):
        #print "start_group %s,%s" % (loc, toks)
        self.push_state()
        return []
    
    def group(self, s, loc, toks):
        #print "group %s,%s" % (loc, toks)
        pass
    
    def end_group(self, s, loc, toks):
        #print "end_group %s,%s" % (loc, toks)
        self.pop_state()
        return []
    
    def font(self, s, loc, toks):
        print "font %s,%s" % (loc, toks)
        
        assert(len(toks)==1)
        
        return []
    
    def is_overunder(self, nucleus):
        print "is-overunder"
    
    def is_dropsub(self, nucleus):
        print "is_dropsub"
        # if isinstance(nucleus, Char):
            # return nucleus.c in self._dropsub_symbols
        return False
    
    def is_slanted(self, nucleus):
        print "is_slanted"
        # if isinstance(nucleus, Char):
            # return nucleus.is_slanted()
        return False
    
    def subsuperscript(self, s, loc, toks):
        #print "subsuperscript %s,%s" % (loc, toks)
        
        nucleus = None
        sub = None
        super = None
        
        if len(toks[0]) == 1:
            return toks[0].asList()
        elif len(toks[0]) == 2:
            op, next = toks[0]
            nucleus = Symbol()
            if op == '_':
                sub = next
            else:
                super = next
        elif len(toks[0]) == 3:
            nucleus, op, next = toks[0]
            if op == '_':
                sub = next
            else:
                super = next
            #print "necleus %s, super %s" % (nucleus,super)
            return [ pow(nucleus,super)]
        elif len(toks[0]) == 5:
            nucleus, op1, next1, op2, next2 = toks[0]
            if op1 == op2:
                if op1 == '_':
                    raise ParseFatalException("Double subscript")
                else:
                    raise ParseFatalException("Double superscript")
            if op1 == '_':
                sub = next1
                super = next2
            else:
                super = next1
                sub = next2
        else:
            raise ParseFatalException(
                "Subscript/superscript sequence is too long. "
                "Use braces { } to remove ambiguity.")
        
    
    def genfrac(self, s, loc, toks):
        print "genfrac %s,%s" % (loc, toks)
        
        assert(len(toks)==1)
        assert(len(toks[0])==6)
        
        # return self._genfrac(*tuple(toks[0]))
    
    def stackrel(self, s, loc, toks):
        print "stackrel %s,%s" % (loc, toks)
        
        assert(len(toks)==1)
        assert(len(toks[0])==2)
        
        num, den = toks[0]
        
        #return []
    
    def auto_sized_delimiter(self, s, loc, toks):
        print "auto_sized_delimiter %s,%s" % (loc, toks)
        
        front, middle, back = toks
        
        # return self._auto_sized_delimiter(front, middle.asList(), back)
    
    def frac(self, s, loc, toks):
        #print "frac %s,%s" % (loc, toks)
        
        assert(len(toks)==1)
        assert(len(toks[0])==2)
        
        num, den = toks[0]
        
        return [ (num[0] / den[0]) ]
    
    def binom(self, s, loc, toks):
        #print "binom %s,%s" % (loc, toks)
        
        assert(len(toks)==1)
        assert(len(toks[0])==2)
        
        n, k = toks[0]
        
        binom_coef = binomial(n[0], k[0])
        return [binom_coef]
    
    def sqrt(self, s, loc, toks):
        #print "sqrt %s,%s" % (loc, toks)
        root, body = toks[0]
        
        if root == None:
            rt = 2
        else:
            rt = root[0]
            
        sqrt_expr = (body[0])**(Rational(1,rt))
        return [sqrt_expr]
    

########################################
# Unit tests
########################################

class TestLaTeXParser(unittest.TestCase):
    
    def setUp(self):
        pass
        #self.teststring = '236'
    
    def testparser(self):
        pass
        #self.assertEqual(ChemSpiderId(self.teststring), self.teststring)

if __name__ == '__main__':
    tests = [
        #r'Differentiate $\frac{\cos(x^3 + x)}{x^2}$ with respect to $x$',
        r'Differentiate $\cos(x)$ with respect to $x$',
        r'Differentiate $\frac{x}{y}$ with respect to $x$',
        r'Differentiate $\sqrt{x^3}$ with respect to $x$',
        r'Integrate $\sin(x)$ with respect to $x$'
        ]
        
    for t in tests:
        print
        print t
        parsedResult = MathtexParser().parse(t)
        print parsedResult[0]
    
    #unittest.main()

